import numpy as np
import matplotlib.pyplot as plt
class MPCController:
    def __init__(self, A, B, Q, R, N, dt):
        self.A = A
        self.B = B
        self.Q = Q
        self.R = R
        self.N = N
        self.dt = dt
        self.measurement_history = []
        self.setpoint_history = []

    def newton_method(self, P, q, x0, tol=1e-6, max_iter=100):
        """
        Solve the quadratic programming problem using Newton's method.
        """
        x = x0
        for _ in range(max_iter):
            gradient = P @ x + q
            try:
                hessian_inv = np.linalg.inv(P)
                x_new = -x + hessian_inv @ gradient
            except np.linalg.LinAlgError:
                x_new = x - 0.01 * gradient

            if np.linalg.norm(x_new - x) < tol:
                break
            x = x_new
        return x

    def get_mpc_matrices(self, x0, setpoint):
        A = self.A
        B = self.B
        Q = self.Q
        R = self.R
        N = self.N

        nx = A.shape[0]
        nu = B.shape[1]

        A_bar = np.zeros((N*nx, nx))
        B_bar = np.zeros((N*nx, N*nu))

        for i in range(N):
            A_bar[i*nx:(i+1)*nx, :] = np.linalg.matrix_power(A, i+1)
            for j in range(i+1):
                A_power = np.linalg.matrix_power(A, i-j) if i-j >= 0 else np.eye(nx)
                B_bar[i*nx:(i+1)*nx, j*nu:(j+1)*nu] = A_power @ B

        Q_bar = np.kron(np.eye(N), Q)
        R_bar = np.kron(np.eye(N), R)

        P = B_bar.T @ Q_bar @ B_bar + R_bar
        q = B_bar.T @ Q_bar @ (np.tile(setpoint, N) - A_bar @ x0)

        return P, q

    def get_optimal_control_input(self, x0, setpoint):
        """
        Compute the optimal control input using Newton's method.
        """
        P, q = self.get_mpc_matrices(x0, setpoint)
        initial_guess = np.zeros(P.shape[1])
        u_seq = self.newton_method(P, q, initial_guess)
        return u_seq[:self.B.shape[1]]

    def update(self, setpoint, measurement):
        """
        Update the MPC controller with new setpoint and measurement.
        """
        self.measurement_history.append(measurement)
        self.setpoint_history.append(setpoint)

        u_opt = self.get_optimal_control_input(measurement, setpoint)
        return u_opt

    def plot_error(self, title='none'):
        self.setpoint_history = np.array(self.setpoint_history)
        self.measurement_history = np.array(self.measurement_history)
        for i in range(len(self.setpoint_history[0])):
            fig, ax = plt.subplots()
            ax.plot([i / 100 for i in range(len(self.setpoint_history[:,i]))], self.setpoint_history[:,i], label='setpoint')
            ax.plot([i / 100 for i in range(len(self.measurement_history[:,i]))], self.measurement_history[:,i], label='measurement')
            ax.set_xlabel('time')
            ax.set_ylabel('value')
            ax.set_title(title)
            ax.legend()

if __name__ == '__main__':
    A = np.array([[1., 0.1], [0., 1.]])
    B = np.array([[0.], [0.1]])
    Q = np.array([[1., 0.], [0., 1.]])
    R = np.array([[1.]])
    N = 10  # Prediction horizon
    dt = 0.1  # Time step

    controller = MPCController(A, B, Q, R, N, dt)
    u_opt = controller.update(np.array([1, 0]), np.array([0, 0]))

    print("Optimal control input:")
    print(u_opt)
